# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/17 14:49
# @Function: 制作核算系统中需要的各个参数tif


# 双线性的插值方法
import glob
import os
from typing import List
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
from interpolate import cut_tif_to_chunks, interpolate_to_target_resolution_01, merge_tif_files_for_one
import xarray as xr

def save_annual_sum_as_geotiff(input_nc_path: str, output_tif_path: str = None, variable_name: str = 'e',
                               time_dim: str = 'valid_time', crs: str = "EPSG:4326"):
    """
    读取NetCDF文件，计算指定变量的年总和，并将其保存为GeoTIFF文件。

    参数:
    input_nc_path (str): 输入NetCDF文件的路径。
    output_tif_path (str): 输出GeoTIFF文件的路径。如果未指定，则默认为输入文件所在文件夹下的'annual_sum.tif'。
    variable_name (str): 要累加的变量名称，默认为'e'。
    time_dim (str): 时间维度的名称，默认为'valid_time'。
    crs (str): 输出文件的坐标参考系统，默认为"EPSG:4326"（WGS84地理坐标系）。
    """
    try:
        # 打开NetCDF数据集
        dataset = xr.open_dataset(input_nc_path)

        # 查看数据集结构
        print(dataset)

        # 对全年数据进行累加
        annual_sum = dataset[variable_name].sum(dim=time_dim)

        # 写入CRS信息（推荐用法）
        annual_sum = annual_sum.rio.write_crs(crs)

        # 设置空间维度（经度和纬度）
        annual_sum.rio.set_spatial_dims(x_dim="longitude", y_dim="latitude", inplace=True)

        # 如果未指定输出路径，则使用默认路径
        if output_tif_path is None:
            input_filename = os.path.basename(input_nc_path)
            input_filename_without_ext = os.path.splitext(input_filename)[0]
            output_tif_path = os.path.join(os.path.dirname(input_nc_path), f'{input_filename_without_ext}_annual_sum.tif')

        # 保存为GeoTIFF
        annual_sum.rio.to_raster(output_tif_path)

        print(f"Annual sum saved to {output_tif_path}")

    except Exception as e:
        print(f"An error occurred: {e}")



# todo 切割的时候将空的crs赋值上默认的坐标系

def process_tif_file(input_tif_path: str, target_resolution: float = 0.00027777778, num_chunks: int = 40) -> str:
    """
    处理输入的 TIFF 文件，包括切割、插值和合并步骤。

    参数:
    input_tif_path (str): 输入的 TIFF 文件路径。
    target_resolution (float): 目标分辨率，默认为 0.007777778。
    num_chunks (int): 切割成的块数，默认为 40。

    返回:
    str: 合并后的 TIFF 文件夹路径。
    """
    # 切割 TIFF 文件
    cut_tif_to_chunks(input_tif_path, num_chunks=num_chunks)

    # 获取切割后的文件夹路径
    input_folder = os.path.dirname(input_tif_path)
    input_filename = os.path.basename(input_tif_path)
    output_folder = os.path.join(input_folder, f"{os.path.splitext(input_filename)[0]}_chunks")

    # 遍历文件夹中的所有 TIFF 文件
    tif_files: List[str] = glob.glob(os.path.join(output_folder, '*.tif'))

    for tif_file in tqdm(tif_files, desc="Processing TIFF files"):
        print(f"Processing file: {tif_file}")
        # 对每个 TIFF 文件进行插值处理
        interpolate_to_target_resolution_01(tif_file, target_resolution=target_resolution)

    # 设置合并后的文件夹路径
    merge_folder = os.path.join(input_folder, f"{os.path.splitext(input_filename)[0]}_interpolation")
    print(f"Merge folder: {merge_folder}")

    # 合并处理后的 TIFF 文件
    merge_tif_files_for_one(merge_folder)

    return merge_folder


def process_tif_file_thread(input_tif_path: str, max_threads:int = 4, target_resolution: float = 0.00027777778, num_chunks: int = 40) -> str:
    """
    处理输入的 TIFF 文件，包括切割、插值和合并步骤。

    参数:
    input_tif_path (str): 输入的 TIFF 文件路径。
    target_resolution (float): 目标分辨率，默认为 0.00027777778。
    num_chunks (int): 切割成的块数，默认为 40。

    返回:
    str: 合并后的 TIFF 文件夹路径。
    """
    # 切割 TIFF 文件
    cut_tif_to_chunks(input_tif_path, num_chunks=num_chunks)

    # 获取切割后的文件夹路径
    input_folder = os.path.dirname(input_tif_path)
    input_filename = os.path.basename(input_tif_path)
    output_folder = os.path.join(input_folder, f"{os.path.splitext(input_filename)[0]}_chunks")

    # 遍历文件夹中的所有 TIFF 文件
    tif_files: List[str] = glob.glob(os.path.join(output_folder, '*.tif'))

    # 指定线程池的最大线程数
    # max_threads = 4  # 线程数量通常不应超过 CPU 核心数

    # 使用多线程处理每个 TIFF 文件的插值操作
    with ThreadPoolExecutor(max_workers=max_threads) as executor:
        list(tqdm(
            executor.map(lambda tif_file: interpolate_to_target_resolution_01(tif_file, target_resolution), tif_files),
            total=len(tif_files), desc="Processing TIFF files"))

    # 设置合并后的文件夹路径
    merge_folder = os.path.join(input_folder, f"{os.path.splitext(input_filename)[0]}_interpolation")
    print(f"Merge folder: {merge_folder}")

    # 合并处理后的 TIFF 文件
    merge_tif_files_for_one(merge_folder)

    return merge_folder


# input_tif_path = r"F:\test_files\插值验证程序的输入数据\双线性插值\et_2023.tif"
# # process_tif_file(input_tif_path)
# process_tif_file_thread(input_tif_path, 4)
